// pubsub_node.cpp
#include <zmq.hpp>
#include <iostream>
#include <string>
#include <thread>
#include <chrono>
#include <atomic>

class PubSubNode
{
  public:
    PubSubNode (const std::string &node_name,
                const std::string &pub_endpoint,
                const std::string &sub_endpoint) :
        node_name_ (node_name),
        pub_endpoint_ (pub_endpoint),
        sub_endpoint_ (sub_endpoint)
    {
    }

    void start ()
    {
        context_ = new zmq::context_t (1);
        pub_socket_ = new zmq::socket_t (*context_, ZMQ_PUB);
        sub_socket_ = new zmq::socket_t (*context_, ZMQ_SUB);

        pub_socket_->connect (pub_endpoint_);
        sub_socket_->connect (sub_endpoint_);

        sub_socket_->setsockopt (ZMQ_SUBSCRIBE, "", 0);

        std::cout << node_name_ << " started" << std::endl;

        // Start the message handling thread
        message_thread_ = std::thread (&PubSubNode::handle_messages, this);
    }

    void subscribe_topic (const std::string &topic)
    {
        sub_socket_->setsockopt (ZMQ_SUBSCRIBE, topic.c_str (),
                                 topic.length ());
        std::cout << node_name_ << " subscribed to topic: " << topic
                  << std::endl;
    }

    void unsubscribe_topic (const std::string &topic)
    {
        sub_socket_->setsockopt (ZMQ_UNSUBSCRIBE, topic.c_str (),
                                 topic.length ());
        std::cout << node_name_ << " unsubscribed from topic: " << topic
                  << std::endl;
    }

    void publish (const std::string &topic, const std::string &message)
    {
        std::string full_message = topic + " " + message;
        zmq::message_t msg (full_message.begin (), full_message.end ());
        pub_socket_->send (msg, zmq::send_flags::none);
        std::cout << node_name_ << " published [" << topic << "]: " << message
                  << std::endl;
    }

    void stop ()
    {
        running_ = false;
        if (message_thread_.joinable ()) {
            message_thread_.join ();
        }
        delete pub_socket_;
        delete sub_socket_;
        delete context_;
    }

    void handle_messages ()
    {
        while (running_) {
            zmq::message_t msg;
            if (sub_socket_->recv (msg, zmq::recv_flags::dontwait)) {
                std::string full_message = msg.to_string ();

                // 解析主题和消息内容
                size_t space_pos = full_message.find (' ');
                if (space_pos != std::string::npos) {
                    std::string topic = full_message.substr (0, space_pos);
                    std::string message = full_message.substr (space_pos + 1);
                    std::cout << node_name_ << " received [" << topic
                              << "]: " << message << std::endl;
                } else {
                    std::cout << node_name_ << " received: " << full_message
                              << std::endl;
                }
            }
            std::this_thread::sleep_for (std::chrono::milliseconds (10));
        }
    }


  private:
    std::string node_name_;
    std::string pub_endpoint_;
    std::string sub_endpoint_;
    zmq::context_t *context_;
    zmq::socket_t *pub_socket_;
    zmq::socket_t *sub_socket_;
    std::thread message_thread_;
    std::atomic<bool> running_{true};
};

int main (int argc, char *argv[])
{
    if (argc != 2) {
        std::cout << "Usage: " << argv[0] << " <node_name(A|B|C)>" << std::endl;
        return 1;
    }
    std::string node_name = argv[1];
    if (node_name != "A" && node_name != "B" && node_name != "C") {
        std::cout << "Node name must be A, B or C" << std::endl;
        return 1;
    }

    PubSubNode node (node_name, "tcp://localhost:5556", "tcp://localhost:5555");
    node.start ();

    // 订阅其他节点的主题消息
    if (node_name == "A") {
        node.subscribe_topic ("B");
        node.subscribe_topic ("C");
    } else if (node_name == "B") {
        node.subscribe_topic ("A");
        node.subscribe_topic ("C");
    } else if (node_name == "C") {
        node.subscribe_topic ("A");
        node.subscribe_topic ("B");
    }

    std::this_thread::sleep_for (
      std::chrono::seconds (1)); // Wait for connections

    if (node_name == "A") {
        node.publish ("A", "msg1");
    } else if (node_name == "B") {
        node.publish ("B", "msg2");
    } else if (node_name == "C") {
        node.publish ("C", "msg3");
    }

    // 取消订阅其他节点的主题消息
    if (node_name == "A") {
        node.unsubscribe_topic  ("B");
        node.unsubscribe_topic  ("C");
    } else if (node_name == "B") {
        node.unsubscribe_topic  ("A");
        node.unsubscribe_topic  ("C");
    } else if (node_name == "C") {
        node.unsubscribe_topic  ("A");
        node.unsubscribe_topic  ("B");
    }
    node.stop ();
    return 0;
}
